import copy

from mitmproxy import http
from ruamel import yaml


class RewriteContent:

    def response(self, flow: http.HTTPFlow):
        if flow.response:
            # 检查响应头部是否存在
            if "x-frame-options" not in flow.response.headers:
                # 如果不存在，添加 X-Frame-Options 头部
                flow.response.headers["X-Frame-Options"] = "ALLOW-FROM http*"
            else:
                # 如果存在，修改 X-Frame-Options 头部的值
                flow.response.headers["X-Frame-Options"] = "ALLOW-FROM http*"
        yml_case_temp = {
            "suite_name": "场景名称_自定义",
            "suite_desc": "场景描述信息_自定义",
            "variables": None,
            "test_steps": []
        }
        yaml_response = copy.deepcopy(flow.response)
        filename = "test_" + flow.request.path.split("?")[0].replace("/", "_").replace("-", "_") + ".yml"
        cd = f"attachment; filename={filename}"
        #
        yaml_response.headers["Content-Disposition"] = cd

        yaml_response.set_content(
            bytes(yaml.round_trip_dump(yml_case_temp, default_flow_style=False, allow_unicode=True), "utf-8"))
        flow.__setattr__("yaml_response", yaml_response)
        # print("yaml_content", flow.yaml_response)


addons = [
    RewriteContent()
]
